#pragma once

#include "kwdb/status.h"
#include <functional>
#include <chrono>
#include <random>
#include <iostream>
#include <thread>

namespace kwdb {

// Retry policy configuration
struct RetryConfig {
    int max_retries = 3;                     // Maximum retry attempts
    int initial_backoff_ms = 100;            // Initial backoff delay
    int max_backoff_ms = 2000;               // Maximum backoff delay
    double backoff_multiplier = 2.0;         // Exponential backoff multiplier
    double jitter_factor = 0.1;              // Random jitter (0.0-1.0)

    // Retryable error codes (gRPC status codes)
    // 14 = UNAVAILABLE, 4 = DEADLINE_EXCEEDED
    std::vector<int> retryable_codes = {14, 4};
};

class RetryPolicy {
public:
    explicit RetryPolicy(const RetryConfig& config = RetryConfig());

    // Execute function with retry logic
    template<typename Func>
    Status Execute(Func func);

    // Check if status code is retryable
    bool IsRetryable(int code) const;

    // Get retry metrics
    size_t GetRetryCount() const { return retry_count_; }
    size_t GetTotalAttempts() const { return total_attempts_; }

private:
    // Calculate next backoff with exponential backoff and jitter
    std::chrono::milliseconds CalculateBackoff(int attempt);

    RetryConfig config_;
    std::mt19937 random_gen_;
    std::uniform_real_distribution<> jitter_dist_;

    // Metrics
    size_t retry_count_ = 0;
    size_t total_attempts_ = 0;
};

// Template implementation
template<typename Func>
Status RetryPolicy::Execute(Func func) {
    Status status;

    for (int attempt = 0; attempt <= config_.max_retries; attempt++) {
        total_attempts_++;

        status = func();

        if (status.ok()) {
            return status;
        }

        // Check if error is retryable
        if (!IsRetryable(status.code) || attempt == config_.max_retries) {
            return status;
        }

        // Calculate backoff
        auto backoff = CalculateBackoff(attempt);

        retry_count_++;
        std::cout << "[RetryPolicy] Retry #" << (attempt + 1)
                  << " after " << backoff.count() << "ms"
                  << " (error: " << status.message << ")" << std::endl;

        // Sleep for backoff duration
        std::this_thread::sleep_for(backoff);
    }

    return status;
}

} // namespace kwdb